package com.campus.counseling.controller;

import com.campus.counseling.entity.TrainingData;
import com.campus.counseling.service.Word2VecService;
import com.campus.counseling.service.impl.PsyQADataLoaderImpl;
import com.campus.counseling.service.impl.TrainingServiceImpl;
import lombok.RequiredArgsConstructor;
import org.springframework.http.ResponseEntity;
import org.springframework.security.access.prepost.PreAuthorize;
import org.springframework.web.bind.annotation.PostMapping;
import org.springframework.web.bind.annotation.RequestBody;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;

import java.io.File;
import java.util.List;
import java.util.stream.Collectors;
import java.util.ArrayList;
import java.nio.file.Files;
import java.nio.charset.StandardCharsets;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

@RestController
@RequestMapping("/api/training")
@RequiredArgsConstructor
public class TrainingController {

    private static final Logger log = LoggerFactory.getLogger(TrainingController.class);

    private final TrainingServiceImpl trainingService;
    private final Word2VecService word2VecService;
    private final PsyQADataLoaderImpl psyQADataLoader;
    
    @PostMapping("/start")
    @PreAuthorize("permitAll()")
    public ResponseEntity<String> startTraining() {
        try {
            trainingService.trainWithPsyQA();
            return ResponseEntity.ok("训练完成");
        } catch (Exception e) {
            return ResponseEntity.badRequest().body("训练失败: " + e.getMessage());
        }
    }
    
    @PostMapping("/word2vec")
    @PreAuthorize("permitAll()")
    public ResponseEntity<String> trainWord2Vec() {
        try {
            List<TrainingData> trainingDataList = psyQADataLoader.loadPsyQAData();
            List<String> texts = trainingDataList.stream()
                .map(data -> data.getQuestion() + " " + data.getDescription())
                .collect(Collectors.toList());
                
            word2VecService.trainModel(texts);
            return ResponseEntity.ok("Word2Vec模型训练完成");
        } catch (Exception e) {
            return ResponseEntity.badRequest().body("Word2Vec训练失败: " + e.getMessage());
        }
    }

    @PostMapping("/word2vec/file")
    @PreAuthorize("permitAll()")
    public ResponseEntity<?> trainWord2VecFromFile(@RequestParam String filePath) {
        try {
            word2VecService.trainFromFile(filePath);
            return ResponseEntity.ok("Word2Vec模型训练成功");
        } catch (Exception e) {
            return ResponseEntity.badRequest().body("训练失败: " + e.getMessage());
        }
    }

    @PostMapping("/word2vec/batch")
    @PreAuthorize("permitAll()")
    public ResponseEntity<String> trainWord2VecFromDirectory(@RequestParam String dirPath) {
        try {
            File dir = new File(dirPath);
            if (!dir.isDirectory()) {
                return ResponseEntity.badRequest().body("指定路径不是目录");
            }

            // 递归获取所有txt文件
            List<String> allLines = new ArrayList<>();
            processDirectory(dir, allLines);
            
            if (allLines.isEmpty()) {
                return ResponseEntity.badRequest().body("目录中没有找到有效的训练数据");
            }

            // 一次性训练所有数据
            word2VecService.trainModel(allLines);
            
            return ResponseEntity.ok(String.format(
                "Word2Vec模型训练完成，共处理 %d 行数据", allLines.size()));
        } catch (Exception e) {
            log.error("训练失败", e);
            return ResponseEntity.badRequest().body("训练失败: " + e.getMessage());
        }
    }

    private void processDirectory(File dir, List<String> allLines) {
        File[] files = dir.listFiles();
        if (files != null) {
            for (File file : files) {
                if (file.isDirectory()) {
                    processDirectory(file, allLines);
                } else if (file.getName().endsWith(".txt")) {
                    try {
                        log.info("正在处理文件: {}", file.getAbsolutePath());
                        List<String> lines = Files.readAllLines(
                            file.toPath(), StandardCharsets.UTF_8);
                        allLines.addAll(lines);
                        log.info("文件处理完成，读取 {} 行", lines.size());
                    } catch (Exception e) {
                        log.error("处理文件失败: " + file.getName(), e);
                    }
                }
            }
        }
    }
} 